import random
from math import gcd


def isprime(num):
    if num < 2:
        return False
    if num <= 3:
        return True
    if num % 2 == 0 or num % 3 == 0:
        return False
    i = 5
    while i * i <= num:
        if num % i == 0 or num % (i + 2) == 0:
            return False
        i += 6
    return True


# 找到两个大素数
def generate_large_prime(keysize=1024):
    while True:
        num = random.randrange(2 ** (keysize - 1), 2 ** keysize)
        if isprime(num):
            return num


# 计算模反元素（扩展欧几里得算法）
def modinv(a, m):
    def egcd(a, b):
        if a == 0:
            return b, 0, 1
        g, x, y = egcd(b % a, a)
        return g, y - (b // a) * x, x

    g, x, y = egcd(a, m)
    if g != 1:
        raise Exception('模反元素不存在')
    return x % m


# 生成RSA密钥对
def generate_keypair(keysize=1024):
    p = generate_large_prime(keysize // 2)
    q = generate_large_prime(keysize // 2)
    print(p, q)
    n = p * q
    phi_n = (p - 1) * (q - 1)

    e = 65537  # 通常选择65537
    d = modinv(e, phi_n)

    return ((n, e), (n, d))


# 加密
def encrypt(public_key, plaintext):
    n, e = public_key
    plaintext_int = int.from_bytes(plaintext.encode('utf-8'), 'big')
    return pow(plaintext_int, e, n)


# 解密
def decrypt(private_key, ciphertext):
    n, d = private_key
    decrypted_int = pow(ciphertext, d, n)
    return decrypted_int.to_bytes((decrypted_int.bit_length() + 7) // 8, 'big').decode('utf-8')


# 测试RSA算法
public_key, private_key = generate_keypair(100)
print("公钥:", public_key)
print("私钥:", private_key)
message = "RSA加密示例"
print("原始消息:", message)

ciphertext = encrypt(public_key, message)
print("加密后的密文:", ciphertext)

decrypted_message = decrypt(private_key, ciphertext)
print("解密后的消息:", decrypted_message)
